I think this logic is wrong? Let's try and work through it:

The options for the allocation are:

1. stream ordered: CU_POINTER_ATTRIBUTE_CONTEXT will return CUDA_SUCCESS and put ctx = nullptr
2. not stream ordered: CU_POINTER_ATTRIBUTE_CONTEXT will return CUDA_SUCCESS and put ctx = some_context
3. not a UVM allocation: CU_POINTER_ATTRIBUTE_CONTEXT will return CUDA_ERROR_INVALID_VALUE
4. Some other thing: some other cuda error

So if we hit case 3 or 4, we want to error, if we hit case 2, we want to return the context that the query returned, if we hit case 1 we need to do some more work: I think we can just use CU_POINTER_ATTRIBUTE_DEVICE_POINTER query and then do something based on that (either return a non-null current context or else build a primary context and return that).

So perhaps:

CUcontext ctx;
// either the pointer is valid (CUDA_SUCCESS), or it isn't (some other error)
CUDA_DRIVER_TRY(cudaAPI::instance().PointerGetAttribute(&ctx, CU_POINTER_ATTRIBUTE_CONTEXT, dev_ptr));
if (ctx != nullptr) {
    // not a stream-ordered allocation, this context is fine.
    return ctx;
}
// stream-ordered allocation, let's see if the current device context exists and we can use it
CUDA_DRIVER_TRY(cudaAPI::instance().CtxGetCurrent(&ctx));
if (ctx != nullptr) {
    CUdeviceptr target{};
    auto const err = cudaAPI::instance().PointerGetAttribute(&target, CU_POINTER_ATTRIBUTE_DEVICE_POINTER, dev_ptr);
    if (err == CUDA_SUCCESS && target == dev_ptr) return ctx;
}
// Either no context, or the current context can't access the device pointer?
int ord = get_device_ordinal_from_pointer(devPtr);
_primary_contexts.try_emplate(ord, ord);
return _primary_contexts.at(ord).ctx;

WDYT?

You are right, I messed up the equality checks and we don’t need to check the accessibility of devPtr when CU_POINTER_ATTRIBUTE_CONTEXT succeed. However, we must survive a CUDA_ERROR_INVALID_VALUE error.

Why do we fall through if the result is CUDA_ERROR_INVALID_VALUE? Ah, I guess if we don't have UVM, this might still be a device pointer?

We don't want to crash when getting CUDA_ERROR_INVALID_VALUE so we cannot use CUDA_DRIVER_TRY(). Instead, we want to get the current context (if any).

I think I don't understand this check. https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__UNIFIED.html#group__CUDA__UNIFIED_1g0c28ed0aff848042bc0533110e45820c says:

Returns in *data the device pointer value through which ptr may be accessed by kernels running in the current CUcontext. The type of data must be CUdeviceptr *.

If there exists no device pointer value through which kernels running in the current CUcontext may access ptr then CUDA_ERROR_INVALID_VALUE is returned.

If there is no current CUcontext then CUDA_ERROR_INVALID_CONTEXT is returned.

Except in the exceptional disjoint addressing cases discussed below, the value returned in *data will equal the input value ptr.

So my reading is that if the driver call was successful, then current_ctx_dev_ptr will either be the same as dev_ptr (usual case) or different (but still valid) in the corner cases where dev_ptr was allocated either by cuMemHostRegister or cuMemHostAlloc with CU_MEMHOSTALLOC_WRITECOMBINED.

In either case, the current context is good for accessing the pointer.

So I think:

If dev_ptr != current_ctx_dev_ptr, we cannot access dev_ptr using the current context right?

No, it means that the memory was allocated in one of the special ways listed above. The context is capable of accessing the data, but must do so through current_ctx_dev_ptr (not dev_ptr). Since this function is only asking "what context is good for accessing this pointer?" I think we should always return the ctx here.

We could have CUDA_ERROR_INVALID_VALUE here (the current context is not set up for allowing us to access this pointer). Is that a recoverable error?

Should certainly handle other errors as failure.

If dev_ptr was allocated with cuMemHostAlloc with CU_MEMHOSTALLOC_WRITECOMBINED, and is the "host address". It looks like a UVM address, but it is unmapped on the device, so cannot be used in a kernel, even with this primary context that we just made. In that case, the pointer that is usable on the device is the one obtained from querying CUDA_POINTER_ATTRIBUTE_DEVICE_POINTER.

After setting the primary context, I think if we go round and query CUDA_POINTER_ATTRIBUTE_DEVICE_POINTER again and find that the result is different from the input, then we get an error.

Maybe we should try and clarify some more with the cuda team.